import numpy as np
import rerun as rr
import trimesh
import os
import pinocchio as pin

class RobotRerunVisualizer:
    """
    机器人 Rerun 可视化工具类
    
    提供机器人模型的 Rerun 可视化功能，包括：
    1. 初始化 Rerun 可视化环境
    2. 加载和显示机器人网格模型
    3. 更新机器人关节状态
    4. 显示手腕位姿和目标位姿
    """
    
    def __init__(self, app_name="RobotRerunVisualizer", spawn=True):
        """
        初始化 Rerun 可视化工具
        
        参数:
            app_name: Rerun 应用名称
            spawn: 是否自动启动 Rerun 查看器
        """
        self.initialized = False
        try:
            print(f"正在初始化 Rerun 可视化工具 ({app_name})...")
            try:
                # 尝试使用 rr.is_connected() 方法
                if not rr.is_connected():
                    rr.init(app_name, spawn=spawn)
            except (AttributeError, TypeError):
                # 如果方法不存在，直接初始化
                try:
                    rr.init(app_name, spawn=spawn)
                except Exception:
                    # 如果已经初始化，可能会抛出异常，忽略它
                    pass
            
            # 设置坐标系
            try:
                rr.log("", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True)
            except Exception:
                # 如果旧版本没有这个枚举，尝试其他方式
                try:
                    rr.log("", rr.ViewCoordinates("right_hand_z_up"), static=True)
                except Exception:
                    # 如果仍然失败，忽略坐标系设置
                    pass
            
            self.initialized = True
            print("Rerun 可视化工具初始化成功！")
        except Exception as e:
            print(f"Rerun 可视化工具初始化失败: {e}")
            self.initialized = False
    
    def load_robot_model(self, robot, namespace="robot"):
        """
        加载机器人模型到 Rerun 可视化
        
        参数:
            robot: pinocchio 机器人对象
            namespace: 命名空间前缀
        """
        if not self.initialized:
            print("Rerun 未初始化，无法加载机器人模型")
            return False
        
        try:
            self.robot = robot
            self.namespace = namespace
            self.link2mesh = self._get_link2mesh()
            self._load_visual_mesh()
            return True
        except Exception as e:
            print(f"加载机器人模型失败: {e}")
            return False
    
    def _get_link2mesh(self):
        """获取机器人链接对应的网格模型"""
        link2mesh = {}
        for visual in self.robot.visual_model.geometryObjects:
            try:
                mesh = trimesh.load_mesh(visual.meshPath)
                name = visual.name
                if name.endswith('_0'):
                    name = name[:-2]  # 移除后缀
                mesh.visual = trimesh.visual.ColorVisuals()
                mesh.visual.vertex_colors = visual.meshColor
                link2mesh[name] = mesh
            except Exception as e:
                print(f"加载网格模型失败 {visual.name}: {e}")
        return link2mesh
    
    def _load_visual_mesh(self):
        """加载可视化网格模型到 Rerun"""
        if not hasattr(self, 'robot') or not hasattr(self, 'link2mesh'):
            print("机器人模型或网格数据未加载")
            return
        
        self.robot.framesForwardKinematics(pin.neutral(self.robot.model))
        for visual in self.robot.visual_model.geometryObjects:
            try:
                frame_name = visual.name
                if frame_name.endswith('_0'):
                    frame_name = frame_name[:-2]  # 移除后缀
                
                if frame_name not in self.link2mesh:
                    continue
                    
                mesh = self.link2mesh[frame_name]
                
                # 获取关节和帧的ID
                try:
                    frame_id = self.robot.model.getFrameId(frame_name)
                except Exception:
                    # 如果找不到精确匹配，尝试查找包含该名称的帧
                    for i in range(self.robot.model.nframes):
                        if frame_name in self.robot.model.frames[i].name:
                            frame_id = i
                            break
                    else:
                        continue  # 如果找不到匹配的帧，跳过这个可视对象
                
                parent_joint_id = self.robot.model.frames[frame_id].parent
                parent_joint_name = self.robot.model.names[parent_joint_id]
                
                # 获取变换
                frame_tf = self.robot.data.oMf[frame_id]
                joint_tf = self.robot.data.oMi[parent_joint_id]
                
                # 记录关节变换
                rr.log(f'{self.namespace}/{parent_joint_name}',
                       rr.Transform3D(translation=joint_tf.translation,
                                      mat3x3=joint_tf.rotation,
                                      axis_length=0.01))
                
                # 计算相对变换并应用到网格
                relative_tf = joint_tf.inverse() * frame_tf
                mesh_copy = mesh.copy()
                mesh_copy.apply_transform(relative_tf.homogeneous)
                
                # 记录网格
                rr.log(f'{self.namespace}/{parent_joint_name}/{frame_name}',
                       rr.Mesh3D(
                           vertex_positions=mesh_copy.vertices,
                           triangle_indices=mesh_copy.faces,
                           vertex_normals=mesh_copy.vertex_normals,
                           vertex_colors=mesh_copy.visual.vertex_colors,
                           albedo_texture=None,
                           vertex_texcoords=None,
                       ),
                       static=True)
            except Exception as e:
                print(f"可视化网格加载失败 {visual.name}: {e}")
    
    def update_robot_state(self, q):
        """
        更新机器人状态
        
        参数:
            q: 关节角度配置
        """
        if not self.initialized or not hasattr(self, 'robot'):
            print("Rerun 未初始化或机器人模型未加载")
            return
            
        # 更新机器人状态
        self.robot.framesForwardKinematics(q)
        
        # 更新可视化
        for visual in self.robot.visual_model.geometryObjects:
            try:
                frame_name = visual.name
                if frame_name.endswith('_0'):
                    frame_name = frame_name[:-2]  # 移除后缀
                
                # 获取关节和帧的ID
                try:
                    frame_id = self.robot.model.getFrameId(frame_name)
                except Exception:
                    # 如果找不到精确匹配，尝试查找包含该名称的帧
                    for i in range(self.robot.model.nframes):
                        if frame_name in self.robot.model.frames[i].name:
                            frame_id = i
                            break
                    else:
                        continue  # 如果找不到匹配的帧，跳过这个可视对象
                
                parent_joint_id = self.robot.model.frames[frame_id].parentJoint
                parent_joint_name = self.robot.model.names[parent_joint_id]
                
                # 获取关节变换
                joint_tf = self.robot.data.oMi[parent_joint_id]
                
                # 更新关节变换
                rr.log(f'{self.namespace}/{parent_joint_name}',
                       rr.Transform3D(translation=joint_tf.translation,
                                      mat3x3=joint_tf.rotation,
                                      axis_length=0.01))
            except Exception:
                pass  # 忽略错误，继续处理其他可视对象
    
    def visualize_wrist_poses(self, left_target=None, right_target=None, 
                              left_current=None, right_current=None, 
                              namespace="wrists", axis_length=0.05):
        """
        可视化手腕位姿
        
        参数:
            left_target: 左手腕目标位姿 (4x4 齐次变换矩阵)
            right_target: 右手腕目标位姿 (4x4 齐次变换矩阵)
            left_current: 左手腕当前位姿 (pinocchio SE3 对象)
            right_current: 右手腕当前位姿 (pinocchio SE3 对象)
            namespace: 命名空间前缀
            axis_length: 坐标轴长度
        """
        if not self.initialized:
            print("Rerun 未初始化，无法可视化手腕位姿")
            return
        
        # 记录左手腕目标位姿
        if left_target is not None:
            rr.log(f"{namespace}/left_target", 
                   rr.Transform3D(translation=left_target[:3, 3], 
                                  mat3x3=left_target[:3, :3],
                                  axis_length=axis_length))
        
        # 记录右手腕目标位姿
        if right_target is not None:
            rr.log(f"{namespace}/right_target", 
                   rr.Transform3D(translation=right_target[:3, 3], 
                                  mat3x3=right_target[:3, :3],
                                  axis_length=axis_length))
        
        # 记录左手腕当前位姿
        if left_current is not None:
            rr.log(f"{namespace}/left_current", 
                   rr.Transform3D(translation=left_current.translation, 
                                  mat3x3=left_current.rotation,
                                  axis_length=axis_length))
        
        # 记录右手腕当前位姿
        if right_current is not None:
            rr.log(f"{namespace}/right_current", 
                   rr.Transform3D(translation=right_current.translation, 
                                  mat3x3=right_current.rotation,
                                  axis_length=axis_length))
    
    def visualize_points(self, points, namespace="points", colors=None, point_size=5.0):
        """
        可视化点云
        
        参数:
            points: 点云数据 (Nx3 数组)
            namespace: 命名空间前缀
            colors: 点云颜色 (Nx3 数组，RGB值范围0-1)
            point_size: 点的大小
        """
        if not self.initialized:
            print("Rerun 未初始化，无法可视化点云")
            return
        
        if colors is None:
            # 默认使用蓝色
            colors = np.ones((len(points), 3)) * np.array([0.0, 0.5, 1.0])
        
        rr.log(namespace,
               rr.Points3D(
                   positions=points,
                   colors=colors,
                   sizes=np.ones(len(points)) * point_size
               ))
    
    def visualize_trajectory(self, points, namespace="trajectory", color=None, radius=0.005):
        """
        可视化轨迹
        
        参数:
            points: 轨迹点 (Nx3 数组)
            namespace: 命名空间前缀
            color: 轨迹颜色 (RGB值范围0-1)
            radius: 轨迹线宽
        """
        if not self.initialized:
            print("Rerun 未初始化，无法可视化轨迹")
            return
        
        if color is None:
            # 默认使用绿色
            color = [0.0, 1.0, 0.5]
        
        rr.log(namespace,
               rr.LineStrip3D(
                   points=points,
                   color=color,
                   radius=radius
               ))
    
    def log_scalar(self, value, namespace="scalar"):
        """
        记录标量值
        
        参数:
            value: 标量值
            namespace: 命名空间前缀
        """
        if not self.initialized:
            print("Rerun 未初始化，无法记录标量值")
            return
        
        rr.log(namespace, rr.Scalar(value))
    
    def log_text(self, text, namespace="text"):
        """
        记录文本
        
        参数:
            text: 文本内容
            namespace: 命名空间前缀
        """
        if not self.initialized:
            print("Rerun 未初始化，无法记录文本")
            return
        
        rr.log(namespace, rr.TextDocument(text))


# 示例用法
if __name__ == "__main__":
    import argparse
    import time
    
    parser = argparse.ArgumentParser(description='机器人 Rerun 可视化工具测试')
    args = parser.parse_args()
    
    # 创建可视化工具实例
    visualizer = RobotRerunVisualizer(app_name="RobotRerunTest")
    
    # 创建一些示例数据
    # 创建一个轨迹
    t = np.linspace(0, 2*np.pi, 100)
    x = np.cos(t)
    y = np.sin(t)
    z = t / (2*np.pi)
    trajectory = np.column_stack((x, y, z))
    
    # 可视化轨迹
    visualizer.visualize_trajectory(trajectory, namespace="demo/trajectory")
    
    # 创建一些随机点
    points = np.random.rand(50, 3) * 2 - 1
    colors = np.random.rand(50, 3)
    visualizer.visualize_points(points, namespace="demo/points", colors=colors)
    
    # 创建一些示例位姿
    left_target = np.eye(4)
    left_target[:3, 3] = [1, 0.5, 0]
    
    right_target = np.eye(4)
    right_target[:3, 3] = [1, -0.5, 0]
    
    # 可视化位姿
    visualizer.visualize_wrist_poses(left_target=left_target, right_target=right_target, 
                                     namespace="demo/wrists")
    
    # 记录一些标量和文本
    for i in range(10):
        visualizer.log_scalar(np.sin(i/5.0), namespace=f"demo/scalars/sin")
        visualizer.log_scalar(np.cos(i/5.0), namespace=f"demo/scalars/cos")
        time.sleep(0.1)
    
    visualizer.log_text("这是一个示例文本", namespace="demo/text")
    
    print("Rerun 可视化示例完成，请在 Rerun 查看器中查看结果")
    input("按回车键退出...")
